#!/usr/bin/env python
#~ Copyright 2008 Brian Simons
#~ 
#~ This file is part of usetv.
#~ 
#~ usetv is free software: you can redistribute it and/or modify
#~ it under the terms of the GNU General Public License as published by
#~ the Free Software Foundation, either version 3 of the License, or
#~ (at your option) any later version.
#~ 
#~ usetv is distributed in the hope that it will be useful,
#~ but WITHOUT ANY WARRANTY; without even the implied warranty of
#~ MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#~ GNU General Public License for more details.
#~ 
#~ You should have received a copy of the GNU General Public License
#~ along with this program.  If not, see <http://www.gnu.org/licenses/>

import feedparser
import logging
import os
import re
import urllib
from datetime import date
from xml.dom import minidom


logging.basicConfig(level=logging.DEBUG,
                format='%(asctime)s %(name)-12s %(levelname)-8s %(message)s',
                datefmt='%m-%d %H:%M',
                filename=os.path.join(os.getcwd(), 'usetv.log'),
                filemode='w')
log = logging.getLogger('UseTV')


class UseTVError(Exception): pass
class NoEpisodesError(UseTVError):pass
class TVRageQuickInfoError(UseTVError):pass


class UseTV:
    feeds = dict()
    
    def __init__(self):
        pass
        
    def get_latest_episode_nzbindex(self, show):
        """
        Return the latest episode of a program
        
        The returned object is guaranteed to have these attributes:
        title -- the name of the episode
        season -- the season number
        episode -- the episode number
        
        """
        log.debug('get_latest_episode_nzbindex')
        # Download and parse the feed
        post_age = '7'
        search_group = ''
        min_collection_size = '150'
        max_collection_size = '0'
        feed_params = urllib.urlencode({'searchitem': show['Name'],
                                        'x': 27,
                                        'y': 10,
                                        'age': post_age,
                                        'group': search_group,
                                        'min_size': min_collection_size,
                                        'max_size': max_collection_size,
                                        'poster': ''})
        key = 'nzbindex ' + show['Name']
        if not self.feeds.has_key(key):
            log.debug('"%s" is not in the feeds cache. Downloading feed'
                      % key)
            feed_url = 'http://www.nzbindex.nl/rss/?searchitem='
            self.feeds[key] = feedparser.parse(feed_url + feed_params)
        return self._get_latest_episode(key, show)
        
    def get_latest_episode_tvnzb(self, show):
        """
        Return the latest episode of a program
        
        The returned object is guaranteed to have these attributes:
        title -- the name of the episode
        season -- the season number
        episode -- the episode number
        
        """
        log.debug('get_latest_episode_tvnzb')
        # Download and parse the feed if it hasn't already been done
        # TVNZB only has one feed for all new shows, so there's no
        # point downloading it more than once per session
        if not self.feeds.has_key('tvnzb'):
            log.debug('"tvnzb" is not in the feeds cache. Downloading feed')
            feed_url = 'http://www.tvnzb.com/tvnzb_new.rss'
            self.feeds['tvnzb'] = feedparser.parse(feed_url)
        return self._get_latest_episode('tvnzb', show)
        
    def _get_latest_episode(self, site_key, show):
        """
        Grab regex matches for episodes of show['Name'] from site and 
        return an object describing the latest episode
        
        """
        log.debug('_get_latest_episode')
        if show.has_key('Episode'):
            log.debug('using episode-based regex')
            exp = re.compile(r'''
                (?:             # do not return this group
                %s              # search on the name of the show
                                # replacing spaces with non-word chars
                .*?)            # followed by anything, non-greedy
                [s]?            # an 's' occurring at most once
                \D              # a non-digit char
                (?P<season>     # begin a named group
                \d{1,2})        # 1-2 digits
                [ex]            # an 'e' or an 'x'
                (?P<episode>    # begin a named group
                \d{1,2})        # 1-2 digits
                \D              # a non-digit char
                ''' % show['Name'].replace(' ','\W'), 
                re.IGNORECASE + re.VERBOSE)
        elif show.has_key('Day'):
            log.debug('using date-based regex')
            exp = re.compile(r'''
                (?:             # the show name gets inserted here, followed
                %s              # by any number of chars non-greedily. We
                                # already know it, so it's not returned
                .*?)            #
                (?P<year_is_first>  # 2 extra numbers at the beginning of the
                \d{2,2})?       # date tell us that the year comes first.
                (?P<date1>      
                \d{2,2})        
                [.]
                (?P<date2>
                \d{2,2})
                [.]
                (?:             # this gets rid of the first two year numbers
                \d{2,2})?       # if the year is last
                (?P<date3>
                \d{2,2})
                ''' % show['Name'].replace(' ','\W'), 
                re.IGNORECASE + re.VERBOSE)
        else:
            log.debug('A malformed show has been passed here')
            raise UseTVError('A malformed show has been passed here')
        matches = list()
        for entry in self.feeds[site_key].entries:
            match = exp.search(entry.title)
            if match:
                match = match.groupdict()
                if match.has_key('season'):
                    log.debug('assigning season and episode to %s' \
                              % entry.title)
                    entry.season = match['season']
                    entry.episode = match['episode']
                else:
                    if match['year_is_first'] is None:
                        log.debug('date format appears to be mm.dd.yy' + \
                                  ' or mm.dd.yyyy in %s' % entry.title)
                        entry.day = match['date2']
                        entry.month = match['date1']
                        entry.year = match['date3']
                    else:
                        log.debug('date format appears to be ' + \
                            'yyyy.mm.dd in %s' % entry.title)
                        entry.day = match['date3']
                        entry.month = match['date2']
                        entry.year = match['date1']
                matches.append(entry)
        # Decide what the most recently aired episode is and return it
        latest_episode = {'season':0, 'episode':0,
                          'day':0, 'month':0, 'year':0}
        for match in matches:
            if match.has_key('episode'):
                if int(match.season) >= int(latest_episode['season']):
                    if int(match.episode) > int(latest_episode['episode']):
                        latest_episode = match
            elif match.has_key('day'):
                if int(match.year) >= int(latest_episode['year']) and \
                   int(match.month) >= int(latest_episode['month']) and \
                   int(match.day) > int(latest_episode['day']):
                    latest_episode = match
        # Finding no episodes raises an exception
        if latest_episode == {'season':0, 'episode':0,
                              'day':0, 'month':0, 'year':0}:
            raise NoEpisodesError('No episodes were found')
        return latest_episode
        
    def get_latest_episode(self, show, functions):
        """
        takes a search string and a tuple of functions as an argument. Runs
        the functions to find the latest episode available. Returns a 
        dictionary-like object with episode information
        
        """
        log.debug('get_latest_episode')
        
        if show.has_key('Season') and show.has_key('Episode'):
            try:
                return self.__get_latest_episode_by_episode(show, functions)
            except NoEpisodesError:
                pass
        elif show.has_key('Year') and show.has_key('Month') and \
           show.has_key('Day'):
            try:
                return self.__get_latest_episode_by_date(show, functions)
            except NoEpisodesError:
                pass
        else:
            log.error('The config for %s does not have a' % show['Name'] +
                      'valid set of parameters and is being skipped')
        
    def __get_latest_episode_by_episode(self, show, functions):
        log.debug('__get_latest_episode_by_episode')
        # Decide what the most recently aired episode is and return it
        latest_episode = {'season':0,'episode':0}
        info = self.get_show_info(show)
        for function in functions:
            try:
                f = function(show)
                if int(latest_episode['season']) <= int(f['season']):
                    if int(latest_episode['episode']) < int(f['episode']):
                        latest_episode = f
                        # Stop searching if TVRage says this is the latest
                        # episode
                        if info['Latest Episode Season'] == f['season'] and \
                           info['Latest Episode Number'] == f['episode']:
                            break
            except NoEpisodesError:
                pass
        # Finding no episodes raises an exception
        if latest_episode == {'season':0,'episode':0}:
            log.debug('episode-based search for %s fruitless' % show['Name'])
            raise NoEpisodesError('No episodes were found')
        return latest_episode

    def __get_latest_episode_by_date(self, show, functions):
        log.debug('__get_latest_episode_by_date')
        # Decide what the most recently aired episode is and return it
        latest_episode = {'year':0, 'month':0, 'day':0}
        for function in functions:
            try:
                f = function(show)
                if int(latest_episode['year']) <= int(f['year']):
                    if int(latest_episode['month']) <= int(f['month']):
                        if int(latest_episode['day']) < int(f['day']):
                            latest_episode = f
            except NoEpisodesError:
                pass
        # Finding no episodes raises an exception
        if latest_episode == {'year':0, 'month':0, 'day':0}:
            log.debug('date-based search for %s fruitless' % show['Name'])
            raise NoEpisodesError('No episodes were found')
        return latest_episode
        
    def grab_nzb(self, episode):
        """Download an episode and save it to disk.
        
        Keyword arguments:
        episode -- a feedpaser.FeedDict object representing an episode
        # url -- The URL of the file to download
        # file_name -- The name that should be given to the file
        
        """
        log.debug('grab_nzb')
        path = os.path.join(os.getcwd(), 'NZB')
        if not os.path.exists(path):
            log.warning('%s directory does not exist. creating it now' % path)
            os.mkdir(path)
        # Decide what the episode's file name should be
        exp = re.compile(r'''
            (?:             # don't return quotes surrounding the file name
            ["])
            (?P<file_name> 
            .+?             # file name is everything between the quotes
            (?:
            \.nzb|\.nfo))   # it should end with .nzb or .nfo, else no match.
                            # notice that the file_name group ends after this
                            # suffix check
            (?:
            ["])
            ''', re.IGNORECASE + re.VERBOSE)
        match = exp.search(episode.title)
        # No matches were found with the above regex
        # we'll proceed with an ugly file name
        if match is None:
            if not episode.title.endswith('.nzb'):
                episode.file_name = episode.title + '.nzb'
            else:
                episode.file_name = episode.title
        # We found a match using the above regex.
        # We can make a pretty file name
        else:
            episode.file_name = \
                match.group('file_name').replace('.nfo','.nzb')
        # Retrieve the file
        log.info('saving episode available at "%s"' % episode.link +
                 'to "%s"' %
                 os.path.join(path, episode.file_name.replace(os.sep, '.')))
        urllib.urlretrieve(episode.link,
                 os.path.join(path, episode.file_name.replace(os.sep, '.')))
        
    def get_show_info(self, show):
        """
        Return a dictionary with info from TVRage about the series specified 
        in show['Name'] argument
        
        The returned dictionary is guaranteed to have these keys:
        Latest Episode Number
        Latest Episode Title
        Latest Episode Air Date
        Latest Episode Season
        
        """
        log.debug('get_show_info')
        key = 'tvrage ' + show['Name']
        if not self.feeds.has_key(key):
            log.debug('"%s" is not in the feeds cache. Downloading feed'
                      % key)
            url = 'http://www.tvrage.com/quickinfo.php?show=' \
                  + urllib.quote(show['Name'])
            # read data from url
            sock = urllib.urlopen(url)
            data = sock.read()
            sock.close()
            # Parse info into a dict
            data = data.split('\n')
            info = dict()
            for element in data:
                list = element.split('@')
                # Check that data format hasn't changed
                if not list[0] is '':
                    if len(list) is not 2 or info.has_key(list[0]):
                        raise TVRageQuickInfoError(
                             'TVRage probably changed their quickinfo format')
                    info[list[0]] = list[1]
            # info['Latest Episode'] should now have the format:
            # '10x20^Unending^Jun/22/2007'
            # This isn't terribly useful yet. First we split it on '^'
            # and add back into the dictionary. Then we split the episode
            # number on 'x' and do the same
            (info['Latest Episode Number'],
               info['Latest Episode Title'],
               info['Latest Episode Air Date']) = \
                                             info['Latest Episode'].split('^')
            (info['Latest Episode Season'],
               info['Latest Episode Number']) = \
                                      info['Latest Episode Number'].split('x')
            self.feeds[key] = info
        return self.feeds[key]
            
    def parse_shows(self, path = os.path.join(os.getcwd(), 'usetv.shows')):
        """
        Parse the XML config file and return a list of dictionaries 
        containing episode information
        
        Accepts an alternate path as an argument. This argument should only 
        be needed for testing
        
        """
        log.debug('parse_shows')
        dom = minidom.parse(path)
        shows = list()
        # iterate through Shows
        for n in dom.firstChild.childNodes:
            # Ignore elements representing carriage returns
            # (which have nodeType == 3)
            if n.nodeType == 1:
                show = dict()
                #Iterate through children of Show
                for m in n.childNodes:
                    # Ignore elements representing carriage returns
                    # (which have nodeType == 3)
                    if m.nodeType == 1:
                        show[m.nodeName] = m.firstChild.data
                shows.append(show)
        return shows
        
    def unparse_shows(self, shows,
                      path = os.path.join(os.getcwd(), 'usetv.shows')):
        """
        Save the XML config file to disk with updated info about the last
        downloaded show
        
        Accepts an alternate path as an argument. This argument should only 
        be needed for testing
        
        """
        log.debug('unparse_shows')
        dom = minidom.parse(path)
        # iterate through Shows
        i = 0
        for n in dom.firstChild.childNodes:
            # Ignore elements representing carriage returns
            # (which have nodeType == 3)
            if n.nodeType == 1:
                #Iterate through children of Show
                for m in n.childNodes:
                    # Ignore elements representing carriage returns
                    # (which have nodeType == 3)
                    if m.nodeType == 1:
                        log.debug('writing %s over %s in config' %
                                  (shows[i][m.nodeName], m.firstChild.data))
                        m.firstChild.data = shows[i][m.nodeName]
                i += 1
        file = open(path, 'w')
        file.write(dom.saveXML(dom.firstChild))
        file.close()
        
    def run(self):
        log.debug('run')
        shows = self.parse_shows()
        for show in shows:
            latest_episode = self.get_latest_episode(show,
                                (self.get_latest_episode_tvnzb,
                                 self.get_latest_episode_nzbindex))
            if latest_episode is None:
                log.info('no new episodes of "%s" found' % show['Name'])            
            elif latest_episode.has_key('day'):
                date1 = date(int(latest_episode.year),
                             int(latest_episode.month),
                             int(latest_episode.day))
                date2 = date(int(show['Year']),
                             int(show['Month']),
                             int(show['Day']))
                if date1 > date2:
                    log.debug('date %s is greater than %s' % (date1, date2))
                    self.grab_nzb(latest_episode)
                    show['Year'] = latest_episode.year
                    show['Month'] = latest_episode.month
                    show['Day'] = latest_episode.day
                else:
                    log.info('"%s" episode %s/%s/%s' %
                             (show['Name'],
                              latest_episode.month,
                              latest_episode.day,
                              latest_episode.year)
                             + ' is not newer than the last one we got: '
                             + '%s/%s/%s' % 
                             (show['Month'],
                              show['Day'],
                              show['Year']))
            elif latest_episode.has_key('episode'):
                if (latest_episode.season >= show['Season'] and
                   latest_episode.episode > show['Episode']) or \
                   (latest_episode.season > show['Season'] and
                   latest_episode.episode == show['Episode']):
                    self.grab_nzb(latest_episode)
                    show['Season'] = latest_episode.season
                    show['Episode'] = latest_episode.episode
                else:
                    log.info('"%s" episode %sx%s' %
                             (show['Name'],
                              latest_episode.season,
                              latest_episode.episode)
                             + ' is not newer than the last one we got: '
                             + '%sx%s' % (show['Season'],show['Episode']))
            else:
                raise UseTVError('malformed show')
        self.unparse_shows(shows)

if __name__ == '__main__':
    d = UseTV()
    d.run()
